#!/usr/bin/env python
'''
Created on 2/11/15

@author: Rich Plevin (rich@plevin.com)
'''

# Copyright (c) 2015, Richard Plevin.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the
#    distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import sys
import argparse
import pandas as pd

PROGRAM = os.path.basename(__file__)
VERSION = "0.4"


def parseArgs():
    parser = argparse.ArgumentParser(description='''Compute the differences in the timeseries data in two or
    more CSV files. The "year columns" in reference case results (the first file given) are subtracted from
    those in each of the other result files. If all rows have the same Units, an additional "global" row is added
    containing the sum of all the other rows. Each comparison is labeled and written to the output CSV file given.
    As a convenience, if any of the filenames is given without an extension, the ".csv" extension is added.
    ''')

    #
    # Required arguments
    #
    parser.add_argument('referenceFile', nargs=1,
                        help='''The name of the reference file whose timeseries data is subtracted
                        from that of each resultFile given.''')

    parser.add_argument('otherFiles', nargs='+', metavar='otherFile',
                        help='''One or more CSV files holding XML query results. The "referenceFile"
                        results are subtracted from each "otherFile" and saved to the named output
                        CSV file. If needed, ".csv" suffixes are added to all files.''')

    parser.add_argument('-o', '--outFile', default='differences.csv',
                        help='''The name of the ".csv" or ".xlsx" file containing the differences
                        between each scenario and the reference. Default is "differences.csv".''')

    parser.add_argument('-c', '--convertOnly', default=False, action="store_true",
                        help='''Convert the given CSV files into an Excel workbook, one sheet per CSV file.''')

    parser.add_argument('-p', '--percentage', default=False, action="store_true",
                        help='''Compute the difference on a percentage basis, i.e.,
                                (X-reference)/reference.''')

    parser.add_argument('-s', '--skiprows', type=int, default=1,
                        help='The number of rows to skip. Default is 1, which works for GCAM batch query output. Use -s0 for outFile.csv')

    parser.add_argument('-V', '--version', action='version', version='%(prog)s ' + VERSION)

    args = parser.parse_args()

    return args


def readGcamCsv(filename, skiprows=1):
    df = pd.read_table(filename, sep=',', skiprows=skiprows, index_col=None)
    return df


def dropExtraCols(df, inplace=True):
    columns = df.columns
    unnamed = 'Unnamed:'    # extra (empty) columns can sneak in; we eliminate them
    dropCols = filter(lambda s: s[0:len(unnamed)] == unnamed, columns)

    unneeded  = set(['scenario', 'Notes', 'Date'])
    columnSet = set(columns)
    dropCols += columnSet & unneeded    # drop any columns in both sets

    resultDF = df.drop(dropCols, axis=1, inplace=inplace)
    return resultDF


def computeDifference(df1, df2, percentage=False):
    '''
    Compute the difference of of the contents of filename1 minus contents of filename2,
    returning a new DataFrame with the difference in all the year columns, plus a new
    row with the sum of all the other rows.
    '''
    assert not (df1 is None or df2 is None), \
        "Can't compute difference; one or more DataFrames is None"

    df1 = dropExtraCols(df1, inplace=False)
    dropExtraCols(df2, inplace=True)

    assert set(df1.columns) == set(df2.columns), \
        "Can't compute difference because result sets have different columns. df1:%s, df2:%s" % (df1.columns, df2.columns)

    yearCols = filter(str.isdigit, df1.columns)
    nonYearCols = list(set(df1.columns) - set(yearCols))

    df1.set_index(nonYearCols, inplace=True)
    df2.set_index(nonYearCols, inplace=True)

    # Compute difference for timeseries values
    diff = df2 - df1

    if percentage:
        diff = diff / df1
        return diff         # makes no sense to sum in this case

    # Add a row for 'global' that sums all yearCols across all regions if all units match
    # uniqueUnits = df1['Units'].unique()
    #
    # if len(uniqueUnits) == 1:
    #     globalSum = diff.sum()
    #     for col in nonYearCols:
    #         globalSum[col] = ''
    #     globalSum['region'] = 'global-sum'          # restore all the text-valued columns so "append" words
    #     globalSum['Units']  = uniqueUnits[0]
    #     globalSum = pd.DataFrame(globalSum).T
    #     globalSum.set_index(nonYearCols, inplace=True)
    #     diff = diff.append(globalSum)

    return diff


def csv2xlsx(inFiles, outFile, skiprows=0):
    csvFiles = map(ensureCSV, inFiles)
    # TBD: catch exception on reading bad CSV file; save error and report at the end
    dframes  = map(lambda fname: readGcamCsv(fname, skiprows=skiprows), csvFiles)

    sheetNum = 1
    outFile = ensureXLSX(outFile)
    with pd.ExcelWriter(outFile, engine='xlsxwriter') as writer:
        for df, fname in zip(dframes, csvFiles):
            sheetName = str(sheetNum)
            sheetNum += 1
            df.to_excel(writer, index=None, sheet_name=sheetName, startrow=2, startcol=0)
            worksheet = writer.sheets[sheetName]
            worksheet.write_string(0, 0, fname)


def ensureCSV(filename):
    'Add a .csv extension to a filename, if there is none. Return the filename.'
    mainPart, extension = os.path.splitext(filename)
    if not extension:
        filename = mainPart + '.csv'

    return filename


def ensureXLSX(filename):
    '''Force an extension of ".xlsx" on a filename'''
    xlsx = '.xlsx'

    mainPart, extension = os.path.splitext(filename)
    if not extension:
        filename = mainPart + xlsx
    elif extension != xlsx:
        filename += xlsx

    return filename


def writeDiffsToCSV(outFile, referenceFile, otherFiles, skiprows=1, percentage=False):
    refDF = readGcamCsv(referenceFile, skiprows=skiprows)

    with open(outFile, 'w') as f:
        for otherFile in otherFiles:
            otherFile = ensureCSV(otherFile)   # add csv extension if needed
            otherDF   = readGcamCsv(otherFile, skiprows=skiprows)

            diff = computeDifference(refDF, otherDF, percentage=percentage)

            csvText = diff.to_csv(None)
            label = "[%s] minus [%s]" % (otherFile, referenceFile)
            if percentage:
                label = "(%s)/%s" % (label, referenceFile)
            f.write("%s\n%s\n" % (label, csvText))


def writeDiffsToXLSX(outFile, referenceFile, otherFiles, skiprows=1, percentage=False):
    with pd.ExcelWriter(outFile, engine='xlsxwriter') as writer:
        sheetNum = 1
        refDF = readGcamCsv(referenceFile, skiprows=skiprows)

        for otherFile in otherFiles:
            otherFile = ensureCSV(otherFile)   # add csv extension if needed
            otherDF   = readGcamCsv(otherFile, skiprows=skiprows)

            sheetName = 'Diff%d' % sheetNum
            sheetNum += 1

            diff = computeDifference(refDF, otherDF, percentage=percentage)

            diff.reset_index(inplace=True)      # convert multi-index into regular column values
            diff.to_excel(writer, index=None, sheet_name=sheetName, startrow=2, startcol=0)

            #workbook  = writer.book
            #worksheet = workbook.add_worksheet(sheetName)
            worksheet = writer.sheets[sheetName]
            label     = "[%s] minus [%s]" % (otherFile, referenceFile)
            if percentage:
                label = "(%s)/%s" % (label, referenceFile)
            worksheet.write_string(0, 0, label)

            startRow = diff.shape[0] + 4
            worksheet.write_string(startRow, 0, otherFile)
            startRow += 2
            otherDF.reset_index(inplace=True)
            otherDF.to_excel(writer, index=None, sheet_name=sheetName, startrow=startRow, startcol=0)

        dropExtraCols(refDF, inplace=True)
        refDF.to_excel(writer, index=None, sheet_name='Reference', startrow=0, startcol=0)


def main():
    args = parseArgs()
    referenceFile = ensureCSV(args.referenceFile[0])

    outFile = args.outFile
    root, ext = os.path.splitext(outFile)
    if not ext:
        outFile = ensureCSV(outFile)
        ext = '.csv'

    extensions = ('.csv', '.xlsx')
    if ext not in extensions:
        print "Output file extension must be one of %s" % extensions
        sys.exit(1)

    otherFiles = args.otherFiles
    percentage = args.percentage
    skiprows   = args.skiprows

    if args.convertOnly:
        files = otherFiles
        files.insert(0, referenceFile)
        csv2xlsx(files, outFile, skiprows=skiprows)
        return

    if ext == '.csv':
        writeDiffsToCSV(outFile, referenceFile, otherFiles, skiprows=skiprows, percentage=percentage)
    else:
        writeDiffsToXLSX(outFile, referenceFile, otherFiles, skiprows=skiprows, percentage=percentage)

if __name__ == '__main__':
    main()
    try:
        main()
    except Exception, e:
        print "csvDiff failed: %s" % e
        sys.exit(1)
